{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1] Using AVX512 InstructionSet!\n",
      "num_vectors: 1000000000 dim: 128\n",
      "Reading...\n"
     ]
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n",
      "\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n",
      "\u001b[1;31mClick <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. \n",
      "\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "import SPTAG\n",
    "import numpy as np\n",
    "\n",
    "# 假设 SIFT1M fbin 文件路径\n",
    "fbin_path = '/home/gary/Code/DiskANN/build/data/bigann/bigann_base.bbin'\n",
    "\n",
    "# 读取fbin文件\n",
    "def read_bbin(file_path):\n",
    "    with open(file_path, 'rb') as f:\n",
    "        num_vectors = np.fromfile(f, dtype=np.int32, count=1)[0]  # 读取向量数量\n",
    "        dim = np.fromfile(f, dtype=np.int32, count=1)[0]  # 读取维度\n",
    "        print(\"num_vectors:\",num_vectors,\"dim:\",dim)\n",
    "        print(\"Reading...\")\n",
    "        data = np.fromfile(f, dtype=np.uint8).reshape(num_vectors, dim)  # 读取所有向量并重塑为矩阵\n",
    "        print(\"Read done\")\n",
    "    return data\n",
    "\n",
    "# 加载SIFT1M数据\n",
    "sift1b_data = read_bbin(fbin_path)\n",
    "\n",
    "# 创建SPTAG索引实例\n",
    "index_builder = SPTAG.AnnIndex('KDT', 'UInt8', sift1b_data.shape[1])\n",
    "index_builder.SetBuildParam('NumberOfThreads', '24', \"Index\")  # 设置构建时使用的线程数\n",
    "index_builder.SetBuildParam('Samples', '50000', \"Index\")       # 设置采样数量\n",
    "\n",
    "# 插入数据\n",
    "print(\"Training...\")\n",
    "index_builder.Build(sift1b_data, sift1b_data.shape[0], False)\n",
    "print(\"Training Done...\")\n",
    "\n",
    "# 保存索引\n",
    "index_builder.Save('sift1b_index_py_full_240812')\n",
    "\n",
    "print(\"索引构建完成并已保存。\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
